import heapq

import topo1


def widest_path(graph, source, target):
    """寻找最大瓶颈容量的路径"""
    # 初始化所有链路的已用带宽（惰性初始化）
    for u, v in graph.edges():
        graph[u][v].setdefault('used', 0)  # 自动初始化
    # 输出链路的容量
    print("==============测试===============")
    for u, v in graph.edges():
        print(f"链路 {u}→{v} | 容量: {graph.edges[u, v]['capacity']}")
    print("==============测试===============")

    max_bottleneck = {node: 0 for node in graph.nodes}
    max_bottleneck[source] = float('inf')
    heap = [(-max_bottleneck[source], source)]  # 使用最大堆
    predecessors = {node: None for node in graph.nodes}

    while heap:
        current_bn_neg, u = heapq.heappop(heap)
        current_bn = -current_bn_neg

        if u == target:
            break

        for v in graph.neighbors(u):
            if graph[u][v].get('type') != 'electrical':
                continue
            # 计算剩余带宽
            remaining = graph[u][v]['capacity'] - graph[u][v]['used']
            new_bn = min(current_bn, remaining)
            if new_bn > max_bottleneck[v]:
                max_bottleneck[v] = new_bn
                predecessors[v] = u
                heapq.heappush(heap, (-new_bn, v))

    # 回溯路径
    path = []
    current = target
    while current is not None:
        path.append(current)
        current = predecessors[current]
    path.reverse()

    return path if path[0] == source else None, max_bottleneck[target]

def allocate_paths(graph, sorted_pairs):
    results = []
    for (src, dst), demand in sorted_pairs:
        path, bottleneck = widest_path(graph, src, dst)
        print(f"通信对 ({src}→{dst}) | 最大瓶颈容量: {bottleneck}")
        if path and bottleneck >= demand:
            # 更新链路使用情况
            for u, v in zip(path[:-1], path[1:]):
                graph[u][v]['used'] += demand
            results.append((path, demand / bottleneck))  # 通信时间
        else:
            print(f"无法找到路径：{src}→{dst}")
            results.append(None)  # 无法分配
    return results

# G=topo1.build_hybrid_topology_shared_tor(4)
# sorted_pairs={('ToR1', 'ToR1'): 0.0, ('ToR1', 'ToR2'): 143.82399900906915, ('ToR1', 'ToR3'): 140.50658205663257}
# # sorted(topo1.matrix_to_communication_pairs(topo1.traffic_matrix, node_prefix="ToR"), key=lambda x: x[1], reverse=True)
# # 执行分配
# allocation = allocate_paths(G,sorted_pairs )

# for i, ((src, dst), _) in enumerate(sorted_pairs):
#     res = allocation[i]
#     if res:
#         path, delay = res
#         print(f"通信对 ({src}→{dst}) | 分配路径: {'→'.join(path)}")
#         print(f"通信时间: {delay:.2f} ms\n")
#     else:
#         print(f"通信对 ({src}→{dst}) | 无法分配路径\n")